from langchain.prompts import StringPromptTemplate
from langchain_openai import OpenAI
import inspect


def hello_word():
    print("hello word!")


PROMPT = """\
你是一个非常有经验和天赋的程序员，现在给你如下函数名称，你会按照如下格式，输出这段代码的名称、源代码、中文解释。
函数名称：{function_name}
源代码：
{source_code}
代码解释:
"""


# 获取源代码
def get_source_code(function_name):
    return inspect.getsource(function_name)


# 自定义的模板class

class CustomPrompt(StringPromptTemplate):

    def format(self, **kwargs) -> str:
        # 获取源代码
        source_code = inspect.getsource(kwargs["function_name"])

        # 生成提示词
        prompt = PROMPT.format(function_name=kwargs["function_name"].__name__, source_code=source_code)
        return prompt


a = CustomPrompt(input_variables=['function_name'])
pm = a.format(function_name=hello_word)

llm = OpenAI(
    max_retries=1,
    max_tokens=1024,
    temperature=0,
    openai_api_key="sk-rgymzddvnnrgwawnkevkdzlxkxfvhbviqgibbfgpvsyoqtur",
    openai_api_base="https://api.siliconflow.cn/v1",
    model="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B"
)


answer = llm.invoke(pm)
print(answer)